# -*- coding: utf-8 -*-
from kafka import KafkaConsumer
import argparse
import time
import json

parser = argparse.ArgumentParser()

startTimeDefault = (int(time.time()) + 1) * 1000
print(startTimeDefault)
parser.add_argument('--bootstrap_servers', type=str, required=True)
parser.add_argument('--topic', type=str, required=True)
parser.add_argument('--startTime', type=str, required=False, default=startTimeDefault + 30 * 1000)
parser.add_argument('--endTime', type=str, required=False, default=startTimeDefault + 60 * 1000)
parser.add_argument('--timeKey', type=str, required=False, default="timestamp")
parser.add_argument('--print', type=str, required=False, default="delay")

args = parser.parse_args()
delay_time_list = []

def consumeKafka(bootstrap_servers, topic):
    consumer = KafkaConsumer(bootstrap_servers=bootstrap_servers, auto_offset_reset="latest")
    consumer.subscribe(topic)
    for message in consumer:
        # print("%s:%d:%d:%d: msg=%s" % (message.topic, message.timestamp, message.partition,message.offset, message.value.decode('utf-8')))
        msgObj = json.loads(message.value.decode('utf-8'))
        if args.print == "delay":
            '''
            message: ConsumerRecord(topic='iotVehicleNumFromUsecase', partition=2, offset=385612, timestamp=1675238572695, timestamp_type=0, key=b'BDOF2006112G03', value=b'{"roadId":"142","rscuSn":"BDOF2006112G03"
            ,"trafficVehicleNumber":76.0,"statisticsTime":1675238572675}', headers=[], checksum=None, serialized_key_size=14, serialized_value_size=101, serialized_header_size=-1)
            
            msgObj: {'roadId': '142', 'rscuSn': 'BDOF2006112G03', 'trafficVehicleNumber': 76.0, 'statisticsTime': 1675238572675}
            '''
            delay_time = message.timestamp - msgObj[args.timeKey]
            print("时延:%dms" % (delay_time))
            delay_time_list.append(delay_time)
            # 大于结束时间10s后，自动退出
            if msgObj[args.timeKey] >= args.endTime + 10 * 1000:
                delay_time_avg = sum(delay_time_list)/len(delay_time_list)
                print("平均时延:%dms" % (delay_time_avg))
                break
        else:
            if args.timeKey == 'timeSec':
                if msgObj[args.timeKey] * 1000 >= args.startTime and msgObj[args.timeKey] * 1000 <= args.endTime:
                    print(msgObj)
                continue
            if msgObj[args.timeKey] >= args.startTime and msgObj[args.timeKey] <= args.endTime:
                if msgObj['roadId'] == 81:
                    print(json.dumps(msgObj))
                    continue
            # 大于结束时间30s后，自动退出
            if msgObj[args.timeKey] >= args.endTime + 30 * 1000:
                break
    consumer.close()


if __name__ == "__main__":
    # python3 test_kafka_diff.py --bootstrap_servers "kafka.cluster.local:8397" --topic "rscu_msg" > rscu_msg_delay.log  默认打印时延
    # python3 test_kafka_diff.py --bootstrap_servers "kafka.cluster.local:8397" --topic "rscu_msg" --print "data" > rscu_msg_data.log打印时间段内的数据用作数据量对比
    consumeKafka(args.bootstrap_servers.split(","), args.topic)
